#include <ShlObj_core.h>
#include <initguid.h>
#include <dbt.h>

#include "log.h"
#include "ClipboardShareWindow.h"
#include "DataObject.h"
#include "public.h"
#include "virtiofs.h"

const int OPEN_CLIPBOARD_COUNT = 10;
const UINT WM_CLIPBOARD_PUT_FILES = RegisterWindowMessage(CLIPBOARD_PUT_FILES_MSG);
const UINT WM_DELETE_SHARE_NAME = RegisterWindowMessage(DELETE_SHARE_NAME_MSG);
const UINT WM_SYNC_DATA_TO_HOST = RegisterWindowMessage(SYNC_CLIPBOARD_DATA_TO_HOST_MSG);

ClipboardShareWindow::ClipboardShareWindow(
    std::function<void()> fetchAndSendClipboardData,
    std::function<void(WPARAM, LPARAM)> devChangeNotify,
    std::function<std::string(FileShareList&)> fileShareUpdate,
    std::function<void(const char*, DWORD)> sendShareName)
    : m_fetchAndSendClipboardData(fetchAndSendClipboardData),
      m_devChangeNotify(devChangeNotify),
      m_fileShareUpdate(fileShareUpdate),
      m_sendShareName(sendShareName),
      m_bInitialised(FALSE),
      m_hWnd(nullptr),
      m_hInst(nullptr),
      m_hWindowThread(nullptr),
      m_hCompleteInitCbs(nullptr),
      m_vioFsNotify(nullptr),
      m_portNotify(nullptr),
      m_clipboardUpdateFlag(FALSE)
{}

ClipboardShareWindow::~ClipboardShareWindow()
{
    Destroy();
}

DWORD ClipboardShareWindow::Init()
{
    if (m_bInitialised) {
        return ERROR_ALREADY_INITIALIZED;
    }
    m_hInst = reinterpret_cast<HINSTANCE>(GetModuleHandle(nullptr));

    DWORD err = ERROR_SUCCESS;
    m_hCompleteInitCbs = CreateEvent(nullptr, FALSE, FALSE, nullptr);
    if (!m_hCompleteInitCbs) {
        err = GetLastError();
        LOGE("CreateEvent failed with error: {}", err);
        return err;
    }

    m_hWindowThread = CreateThread(nullptr, 0, reinterpret_cast<LPTHREAD_START_ROUTINE>(WindowThread),
        reinterpret_cast<LPVOID>(this), 0, nullptr);
    if (m_hWindowThread == nullptr) {
        err = GetLastError();
        LOGE("CreateThread failed with error: {}", err);
        CloseHandle(m_hCompleteInitCbs);
        return err;
    }

    HANDLE hWaitObjWindowThread[] = {m_hCompleteInitCbs, m_hWindowThread};
    // 2 Indicates the size of hWaitObjWindowThread.
    if (WaitForMultipleObjects(2, hWaitObjWindowThread, FALSE, INFINITE) != WAIT_OBJECT_0) {
        err = GetExitCodeThread(m_hWindowThread, &err) ? err : GetLastError();
        LOGE("Clipboard Share Window Init Failed!");
        CloseHandle(m_hCompleteInitCbs);
        CloseHandle(m_hWindowThread);
        return err;
    }

    CloseHandle(m_hCompleteInitCbs);
    m_hCompleteInitCbs = nullptr;
    m_bInitialised = TRUE;
    LOGI("Clipboard Share Window Init Completed");
    return err;
}

void ClipboardShareWindow::Destroy()
{
    if (!m_bInitialised) {
        LOGE("Destroy failed: Window is not initialised");
        return;
    }

    if (m_portNotify != nullptr) {
        UnregisterDeviceNotification(m_portNotify);
        m_portNotify = nullptr;
    }

    if (m_vioFsNotify != nullptr) {
        UnregisterDeviceNotification(m_vioFsNotify);
        m_vioFsNotify = nullptr;
    }

    SendMessage(m_hWnd, WM_CLOSE, 0, 0);
    if (WaitForSingleObject(m_hWindowThread, TERM_TIMEOUT) == WAIT_TIMEOUT) {
        LOGW("Terminate Thread for Clipboard Window");
        TerminateThread(m_hWindowThread, 0);
    }
    CloseHandle(m_hWindowThread);
    m_hWindowThread = nullptr;
    m_hWnd = nullptr;
    m_hInst = nullptr;
    m_bInitialised = FALSE;
}

BOOL ClipboardShareWindow::OpenClipboard()
{
    /* "OpenClipboard fails if another window has the clipboard open."
     * So try a few times and wait up to 1 second.
     */
    BOOL fOpened = FALSE;
    DWORD dwErr = ERROR_SUCCESS;
    int i;

    for (i = 0; i < OPEN_CLIPBOARD_COUNT; i++) {
        if (::OpenClipboard(m_hWnd)) {
            fOpened = TRUE;
            break;
        }

        if (i == OPEN_CLIPBOARD_COUNT - 1) {
            dwErr = GetLastError();
        }

        /* sleep interval = [1..512] ms */
        Sleep(1 << i);
    }

    if (i > 0) {
        LOGW("{} times tried to open clipboard", i + 1);
    }

    if (!fOpened) {
        LOGE("Failed to open clipboard with error: {}", dwErr);
    }

    return fOpened;
}

BOOL ClipboardShareWindow::CloseClipboard()
{
    BOOL fClosed = FALSE;

    if (::CloseClipboard() == FALSE) {
        const DWORD dwLastErr = GetLastError();
        if (dwLastErr == ERROR_CLIPBOARD_NOT_OPEN) {
            fClosed = TRUE; /* Not important */
        } else {
            LOGE("Failed with error: {}", dwLastErr);
        }
    } else {
        fClosed = TRUE;
    }

    return fClosed;
}

BOOL ClipboardShareWindow::EmptyClipboard()
{
    if (::EmptyClipboard() == FALSE) {
        LOGE("Empty clipboard failed with error: {}", GetLastError());
        return FALSE;
    }
    return TRUE;
}

DWORD ClipboardShareWindow::ClipboardFormatToCustomFormat(DWORD dwFormat)
{
    /* Insert the requested clipboard format data into the clipboard. */
    DWORD format = CBSH_FMT_NONE;

    switch (dwFormat) {
        case CF_UNICODETEXT:
            format = CBSH_FMT_UNICODETEXT;
            break;
        case CF_DIB:
            format = CBSH_FMT_BITMAP;
            break;
        case CF_HDROP:
            format = CBSH_FMT_URI_LIST;
            break;
        default:
            /** Formats registered with RegisterClipboardFormat() start at this index. */
            if (dwFormat >= 0xC000) {
                WCHAR formatName[256];
                int actual = GetClipboardFormatName(dwFormat, formatName, sizeof(formatName) / sizeof(WCHAR));
                if (actual && wcscmp(formatName, CBSH_WIN_REGFMT_HTML) == 0) {
                    format = CBSH_FMT_HTML;
                }
            }
            break;
    }
    return format;
}

ShareSendProcess *ClipboardShareWindow::GetClipboardShareSendProcess(DWORD dwFormat)
{
    /* Priority of Windows clipboard format sync to Host side: File > Text+HTML > Image. */
    if (dwFormat & CBSH_FMT_URI_LIST) {
        return new ClipboardShareSendFilesProcess(m_fileShareUpdate);
    } else if ((dwFormat & CBSH_FMT_UNICODETEXT) || (dwFormat & CBSH_FMT_HTML)) {
        return new ClipboardShareSendTextProcess(m_fileShareUpdate, m_sendShareName, dwFormat);
    } else if (dwFormat & CBSH_FMT_BITMAP) {
        return new ClipboardShareSendBMPProcess();
    } else {
        return nullptr;
    }
}

BOOL ClipboardShareWindow::StartClipboardListen()
{
    return AddClipboardFormatListener(m_hWnd);
}

BOOL ClipboardShareWindow::StopClipboardListen()
{
    return RemoveClipboardFormatListener(m_hWnd);
}

BOOL ClipboardShareWindow::RegisterWindowClass(PCWSTR pszClassName, PCWSTR pszMenuName, WNDPROC lpfnWndProc)
{
    WNDCLASSEX wcex = { sizeof(wcex) };
    wcex.style = CS_HREDRAW | CS_VREDRAW;
    wcex.lpfnWndProc = lpfnWndProc;
    wcex.hInstance = m_hInst;
    wcex.hIcon = nullptr;
    wcex.hCursor = LoadCursor(nullptr, IDC_ARROW);
    wcex.hbrBackground = reinterpret_cast<HBRUSH>(COLOR_WINDOWFRAME);
    wcex.lpszMenuName = pszMenuName;
    wcex.lpszClassName = pszClassName;
    if (!RegisterClassEx(&wcex)) {
        LOGE("RegisterClassEx failed with error: {}", GetLastError());
        return FALSE;
    }
    return TRUE;
}

LRESULT ClipboardShareWindow::HandleRegisterMessage(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    if (msg == WM_SYNC_DATA_TO_HOST) {
        m_fetchAndSendClipboardData();
        return 0;
    } else if (msg == WM_CLIPBOARD_PUT_FILES) {
        /* The shareName must not be null because it is created using 'new' and
         * checked before calling PostMessage in the PutFilesToClipboard function.
         */
        std::string* shareName = (std::string*)lParam;
        ProcessPutFilesToClipboard(shareName->c_str());
        delete shareName;
        return 0;
    } else if (msg == WM_DELETE_SHARE_NAME) {
        /* The shareName must not be null because it is created using 'new' and
         * checked before calling PostMessage in the DeleteShareName function.
         */
        std::string* shareName = (std::string*)lParam;
        ProcessDeleteShareName(shareName->c_str());
        delete shareName;
        return 0;
    } else {
        return DefWindowProc(hWnd, msg, wParam, lParam);
    }
}

LRESULT CALLBACK ClipboardShareWindow::ClipboardShareWndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    ClipboardShareWindow *cbsWnd = reinterpret_cast<ClipboardShareWindow*>(GetWindowLongPtr(hWnd, GWLP_USERDATA));

    switch (msg) {
        case WM_CLOSE:
            DestroyWindow(hWnd);
            break;
        case WM_DESTROY:
            PostQuitMessage(0);
            break;
        case WM_CLIPBOARDUPDATE:
            if (cbsWnd) {
                cbsWnd->SetClipboardUpdateFlag();
            }
            break;
        case WM_DEVICECHANGE:
            if (cbsWnd) {
                cbsWnd->m_devChangeNotify(wParam, lParam);
            }
            break;
        default:
            if (cbsWnd) {
                return cbsWnd->HandleRegisterMessage(hWnd, msg, wParam, lParam);
            }
            return DefWindowProc(hWnd, msg, wParam, lParam);
    }
    return 0;
}

HDEVNOTIFY ClipboardShareWindow::RegisterInterfaceNotify(GUID InterfaceClassGuid)
{
    DEV_BROADCAST_DEVICEINTERFACE NotificationFilter;
    HDEVNOTIFY Notify;

    ZeroMemory(&NotificationFilter, sizeof(NotificationFilter));
    NotificationFilter.dbcc_size = sizeof(DEV_BROADCAST_DEVICEINTERFACE);
    NotificationFilter.dbcc_devicetype = DBT_DEVTYP_DEVICEINTERFACE;
    NotificationFilter.dbcc_classguid = InterfaceClassGuid;

    Notify = RegisterDeviceNotification(m_hWnd, &NotificationFilter, DEVICE_NOTIFY_WINDOW_HANDLE);
    if (Notify == nullptr) {
        LOGE("RegisterDeviceNotification failed with error: {}", GetLastError());
    }

    return Notify;
}

HDEVNOTIFY ClipboardShareWindow::RegisterHandleNotify(HANDLE handle)
{
    if (handle == INVALID_HANDLE_VALUE) {
        LOGE("Invalid handle value");
        return nullptr;
    }

    DEV_BROADCAST_HANDLE NotificationFilter;
    HDEVNOTIFY Notify;

    ZeroMemory(&NotificationFilter, sizeof(NotificationFilter));
    NotificationFilter.dbch_size = sizeof(DEV_BROADCAST_HANDLE);
    NotificationFilter.dbch_devicetype = DBT_DEVTYP_HANDLE;
    NotificationFilter.dbch_handle = handle;
    NotificationFilter.dbch_eventguid = GUID_VIOSERIAL_PORT_CHANGE_STATUS;

    Notify = RegisterDeviceNotification(m_hWnd, &NotificationFilter, DEVICE_NOTIFY_WINDOW_HANDLE);
    if (Notify == nullptr) {
        LOGE("RegisterDeviceNotification failed: {}", GetLastError());
    }

    return Notify;
}

HANDLE ClipboardShareWindow::GetWaitHandle()
{
    return m_hWindowThread;
}

DWORD WINAPI ClipboardShareWindow::WindowThread(ClipboardShareWindow *ptr)
{
    return ptr->CbsWndThread();
}

DWORD ClipboardShareWindow::CbsWndThread()
{
    if (!RegisterWindowClass(CBS_WIN_CLASS_NAME, nullptr, ClipboardShareWndProc)) {
        DWORD err = GetLastError();
        LOGE("Register Clipboard share Window failed");
        return err;
    }

    m_hWnd = CreateWindow(CBS_WIN_CLASS_NAME, nullptr, WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, 0, CBS_WND_WIDTH,
        CBS_WND_HEIGHT, nullptr, nullptr, m_hInst, this);
    if (m_hWnd == nullptr) {
        DWORD err = GetLastError();
        LOGE("CreateWindow failed with error: {}", err);
        return err;
    }

    HRESULT hr = OleInitialize(nullptr);
    if (hr != S_OK) {
        DestroyWindow(m_hWnd);
        return hr;
    }

    SetWindowLongPtr(m_hWnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(this));
    m_portNotify = RegisterInterfaceNotify(GUID_VIOSERIAL_PORT);
    m_vioFsNotify = RegisterInterfaceNotify(GUID_DEVINTERFACE_VIRT_FS);

    if (!AddClipboardFormatListener(m_hWnd)) {
        DWORD err = GetLastError();
        OleUninitialize();
        DestroyWindow(m_hWnd);
        LOGE("AddClipboardFormatListener failed with error: {}", err);
        return err;
    }

    ShowWindow(m_hWnd, SW_HIDE);

    SetEvent(m_hCompleteInitCbs);

    MSG msg;
    while (GetMessage(&msg, nullptr, 0, 0) > 0) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    RemoveClipboardFormatListener(m_hWnd);
    OleUninitialize();
    return ERROR_SUCCESS;
}

void ClipboardShareWindow::DeleteShareName(LPCSTR share)
{
    if (share == nullptr) {
        LOGE("Share name is a null pointer");
        return;
    }

    std::string* shareName = new std::string(share);
    if (shareName == nullptr) {
        LOGE("Failed to allocate share name memory");
        return;
    }

    PostMessage(m_hWnd, WM_DELETE_SHARE_NAME, NULL, (LPARAM)shareName);
}

void ClipboardShareWindow::ProcessDeleteShareName(LPCSTR share)
{
    if (share == nullptr) {
        LOGE("share name is a null pointer");
        return;
    }

    std::string shareName = share;
    m_sendShareName(shareName.c_str(), CBSH_FILES_FINISH);
}

void ClipboardShareWindow::PutFilesToClipboard(LPCSTR share)
{
    if (share == nullptr) {
        LOGE("Input parameters are invalid");
        return;
    }

    std::string* shareName = new std::string(share);
    if (shareName == nullptr) {
        LOGE("Failed to allocate share name memory");
        return;
    }

    PostMessage(m_hWnd, WM_CLIPBOARD_PUT_FILES, NULL, (LPARAM)shareName);
}

void ClipboardShareWindow::SetClipboardUpdateFlag()
{
    m_clipboardUpdateFlag.store(TRUE);
}

void ClipboardShareWindow::SyncClipboardDataToHost()
{
    BOOL expected = TRUE;
    if (m_clipboardUpdateFlag.compare_exchange_strong(expected, FALSE)) {
        PostMessage(m_hWnd, WM_SYNC_DATA_TO_HOST, NULL, NULL);
    }
}

BOOL ClipboardShareWindow::ProcessPutFilesToClipboard(LPCSTR share)
{
    if (share == nullptr) {
        LOGE("share name is a null pointer");
        return FALSE;
    }

    std::string virtFolderPath = std::string(VIRTIOFSD_TEMP_PATH) + "\\" + share;
    LPITEMIDLIST virtFolderItem = ILCreateFromPathA(virtFolderPath.c_str());
    if (!virtFolderItem) {
        LOGE("ILCreateFromPathA failed");
        return FALSE;
    }

    IShellFolder* shellVirtFolder = nullptr;
    HRESULT hr = SHBindToObject(nullptr, virtFolderItem, nullptr, IID_IShellFolder, (void**)&shellVirtFolder);
    if (FAILED(hr)) {
        LOGE("SHBindToObject failed with error: {}", hr);
        ILFree(virtFolderItem);
        return FALSE;
    }

    IEnumIDList* enumIDList = nullptr;
    std::vector<LPITEMIDLIST> itemList;
    hr = shellVirtFolder->EnumObjects(nullptr,
                                      SHCONTF_FOLDERS | SHCONTF_NONFOLDERS | SHCONTF_INCLUDEHIDDEN,
                                      &enumIDList);
    if (SUCCEEDED(hr)) {
        LPITEMIDLIST item = nullptr;
        while (enumIDList->Next(1, &item, nullptr) == S_OK) {
            itemList.push_back(item);
        }
        enumIDList->Release();
    } else {
        LOGE("EnumObjects failed with error: {}", hr);
        ILFree(virtFolderItem);
        return FALSE;
    }

    IDataObject* obj = nullptr;
    hr = shellVirtFolder->GetUIObjectOf(m_hWnd, (UINT)itemList.size(), (LPCITEMIDLIST*)itemList.data(),
                                        IID_IDataObject, nullptr, (void**)&obj);
    shellVirtFolder->Release();
    for (int i = 0; i < itemList.size(); i++) {
        ILFree(itemList[i]);
    }
    ILFree(virtFolderItem);

    if (FAILED(hr)) {
        LOGE("GetUIObjectOf failed with error: {}", hr);
        return FALSE;
    }

    return SetOleClipboard(obj, share);
}

BOOL ClipboardShareWindow::SetOleClipboard(IDataObject* obj, LPCSTR share)
{
    if (!obj) {
        LOGE("obj is a null pointer");
        return FALSE;
    }

    if (share == nullptr) {
        LOGE("share name is a null pointer");
        return FALSE;
    }

    IDataObject* dataObject = new DataObject(obj, m_sendShareName, share);
    if (dataObject == nullptr) {
        LOGE("Failed to allocate memory for IDataObject");
        return FALSE;
    }

    StopClipboardListen();
    HRESULT ret = OleSetClipboard(dataObject);
    StartClipboardListen();
    dataObject->Release();

    return ret == S_OK;
}
